﻿
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
using OptixDotNet;
using OptixDotNet.Math;
using RayDen.Library.Components.SystemComponents;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Devices;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Data;
using DeviceBuffer = OptixDotNet.Buffer;
using Geometry = OptixDotNet.Geometry;
using RayHit = RayDen.Library.Core.RayHit;

namespace RayDen.RayEngine.IntersectionEngines.Optix
{
    public class OptixIntersectionDevice : IntersectionEngine
    {
        private string IntersecitonProgPath = @"G:\Git\RayDen\RayDen.OptixEngine\Cuda\Scripts\CudaScripts\triangle_mesh_id.ptx";
        private string BoundingBoxProgPath =@"G:\Git\RayDen\RayDen.OptixEngine\Cuda\Scripts\CudaScripts\triangle_mesh_id.ptx";
        private string IntersecitonProgName = "mesh_intersect";
        private string BoundingBoxProgName = "mesh_bounds";

        protected new RayEngineScene scene;
        protected Stopwatch wallclock;
        protected double statsDeviceIdleTime;
        protected double statsDeviceTotalTime;
        protected Task intersectionTask;
        protected CancellationTokenSource cancelToken;

        protected InputRayBufferCollection todoRayBuffers;
        protected OutputRayBufferCollection doneRayBuffers;


        protected DeviceBuffer RayHits;
        protected DeviceBuffer Rays;
        protected GeometryGroup GeoGroup { get; set; }

        protected AccelBuilder Builder { get; set; }
        protected AccelTraverser Traverser { get; set; }
        protected Context Context = null;

        public OptixIntersectionDevice(IRayEngineScene scene, bool lowLatency = false) : base(scene)
        {
            this.lowLatency = lowLatency;
            Context = new Context { CpuNumThreads = 4, RayTypeCount = 1, EntryPointCount = 1 };
            
            var rayBufferDesc = new BufferDesc()
            {
                Format = Format.User,
                Width = (uint)(lowLatency ? RayBuffer.RayBufferSize / 8 : RayBuffer.RayBufferSize),
                Type = BufferType.InputOutput,
                ElemSize = (uint)Marshal.SizeOf(typeof(RayData))
            };
            var rayHitBufferDesc = new BufferDesc()
            {
                Format = Format.User,
                Width =  (uint)(lowLatency ? RayBuffer.RayBufferSize / 8 : RayBuffer.RayBufferSize),
                Type = BufferType.InputOutput,
                ElemSize = (uint)Marshal.SizeOf(typeof(RayHit))
            };

            this.todoRayBuffers = new InputRayBufferCollection();
            this.doneRayBuffers = new OutputRayBufferCollection();
            this.started = false;

            Rays = new DeviceBuffer(Context, rayBufferDesc);
            RayHits = new DeviceBuffer(Context, rayHitBufferDesc);
            Context["rayHits"].Set(RayHits);
            Context["rays"].Set(Rays);
            Builder = AccelBuilder.TriangleKdTree;
            Traverser = AccelTraverser.KdTree;
            GeoGroup = new GeometryGroup(Context);
            this.SetScene(scene);
        }

        public override void SetScene(IRayEngineScene secn)
        {
            string shaderPath = Path.GetFullPath(@"G:\Git\RayDen\RayDen.OptixEngine\Cuda\Scripts\CudaScripts\hitkernel.ptx");

            Tracer.TraceLine("OptixIntersectionDevice  -  Initializing Optix");
            this.scene = (RayEngineScene) secn;
            var scn = scene.SceneGeometry;

            Material material = new Material(Context);
            material.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "first_hit");
            //            material.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "shadow");


            var mVertices = scn.Vertices.ToList();
            var mNormals = scn.Normals.ToList();
            List<UV> mTexcoords = null;
            mTexcoords = scn.TexCoords == null ? new List<UV>() : scn.TexCoords.Select(it=>new UV(it.x,it.y)).ToList();

            var vDesc = new BufferDesc() { Width = (uint)mVertices.Count, Format = Format.Float3, Type = BufferType.Input };
            var nDesc = new BufferDesc() { Width = (uint)mNormals.Count, Format = Format.Float3, Type = BufferType.Input };
            var tcDesc = new BufferDesc() { Width = (uint)mTexcoords.Count, Format = Format.Float2, Type = BufferType.Input };

            // Create the buffers to hold our geometry data
            var vBuffer = new  DeviceBuffer(Context, vDesc);
            var nBuffer = new  DeviceBuffer(Context, nDesc);
            var tcBuffer = new DeviceBuffer(Context, tcDesc);


            vBuffer.SetData(mVertices.ToArray());
            nBuffer.SetData(mNormals.ToArray());
            tcBuffer.SetData(mTexcoords.ToArray());


            foreach (var triangleMesh in scene.Meshes.Cast<TriangleMeshInfo>())
            {
                var group = triangleMesh;


                var viDesc = new BufferDesc() { Width = (uint)(group.TrianglesCount), Format = Format.Int3, Type = BufferType.Input };
                var niDesc = new BufferDesc() { Width = (uint)(group.TrianglesCount), Format = Format.Int3, Type = BufferType.Input };
                var tiDesc = new BufferDesc() { Width = (uint)(group.TrianglesCount), Format = Format.Int3, Type = BufferType.Input };
                var ofsDesc = new BufferDesc()
                    {
                        Width = (uint) (group.TrianglesCount),
                        Format = Format.UInt,
                        Type = BufferType.Input
                    };

                var viBuffer = new DeviceBuffer(Context, viDesc);
                var niBuffer = new DeviceBuffer(Context, niDesc);
                var tiBuffer = new DeviceBuffer(Context, tiDesc);

                var gIndexes =
                    scene.Triangles.ToList().GetRange(group.StartTriangle, group.TrianglesCount)
                         .SelectMany(tri => new []
                             {
                                 tri.v0.VertexIndex,
                                 tri.v1.VertexIndex,
                                 tri.v2.VertexIndex
                             })
                         .ToArray();

                var gNIndexes =
                    scene.Triangles.ToList().GetRange(group.StartTriangle, group.TrianglesCount)
                         .SelectMany(tri => new[]
                             {
                                 tri.v0.NormalIndex,
                                 tri.v1.NormalIndex,
                                 tri.v2.NormalIndex
                             })
                         .ToArray();
                var gTIndexes =
                    scene.Triangles.ToList().GetRange(group.StartTriangle, group.TrianglesCount)
                         .SelectMany(tri => new []
                             {
                                 tri.v0.TexCoordIndex,
                                 tri.v1.TexCoordIndex,
                                 tri.v2.TexCoordIndex
                             })
                         .ToArray();

                viBuffer.SetData(Convert(gIndexes).ToArray());
                niBuffer.SetData(Convert(gNIndexes).ToArray());
                tiBuffer.SetData(Convert(gTIndexes).ToArray());

                var geometry = new Geometry(Context);
                geometry.IntersectionProgram = new Program(Context, IntersecitonProgPath, IntersecitonProgName);
                geometry.BoundingBoxProgram = new Program(Context, BoundingBoxProgPath, BoundingBoxProgName);
                geometry.PrimitiveCount = (uint)(group.EndTriangle - group.StartTriangle);

                geometry["vertex_buffer"].Set(vBuffer);
                geometry["normal_buffer"].Set(nBuffer);
                geometry["texcoord_buffer"].Set(tcBuffer);
                geometry["vindex_buffer"].Set(viBuffer);
                geometry["nindex_buffer"].Set(niBuffer);
                geometry["tindex_buffer"].Set(tiBuffer);
                geometry["mesh_offset"].SetInt4(group.StartTriangle);
                //geometry[ "material_buffer" ].Set(mtBuffer);

                //create a geometry instance
                var instance = new GeometryInstance(Context);
                instance.Geometry = geometry;
                instance.AddMaterial(material);
                //create an acceleration structure for the geometry
                Acceleration accel = new Acceleration(Context, Builder, Traverser);

                if (Builder == AccelBuilder.Sbvh || Builder == AccelBuilder.TriangleKdTree)
                {
                    accel.VertexBufferName = "vertex_buffer";
                    accel.IndexBufferName = "vindex_buffer";
                }

                //now attach the instance and accel to the geometry group
                GeoGroup.Acceleration = accel;
                GeoGroup.AddChildren(new[] { instance });
            }

            Program rayGen = new Program(Context, shaderPath, "intersector_camera");
            Program exception = new Program(Context, shaderPath, "exception");
            Program miss = new Program(Context, shaderPath, "miss");
            Context["intersect_ray_type"].Set(0u);
            Context["scene_epsilon"].Set(0.0001f);
            Context.SetRayGenerationProgram(0, rayGen);
            Context.SetExceptionProgram(0, exception);
            Context.SetRayMissProgram(0, miss);
            Context["top_object"].Set(GeoGroup);
            Tracer.TraceLine("OptixIntersectionDevice  -  Compiling Optix contex");
            Context.Compile();
            Context.BuildAccelTree();
            Tracer.TraceLine("OptixIntersectionDevice  -  Complete");
        }

        private IEnumerable<Int3> Convert(int[] list)
        {
            for (int i = 0; i < list.Length; i += 3)
            {
                yield return new Int3(list[i], list[i + 1], list[i + 2]);
            }
        }

        public override void PushRayBuffer(RayBuffer rayBuffer, int clientIdx)
        {
            if (clientIdx + 1 > doneRayBuffers.Count)
            {
                while (doneRayBuffers.Count < (clientIdx + 1))
                {
                    doneRayBuffers.AddBuffer();
                }
            }
            todoRayBuffers.Put(new Tuple<int, RayBuffer>(clientIdx, rayBuffer));
        }

        public override RayBuffer PopRayBuffer(int info)
        {
            return doneRayBuffers.Get(info);
        }

        public override void Stop()
        {
            wallclock.Stop();
            started = false;
            wallclock.Restart();
            if (intersectionTask.Status == TaskStatus.Running)
            {
                this.cancelToken.Cancel();
                Tracer.TraceLine("Attempting to stop Optix intersection device");
                while ((intersectionTask.Status == TaskStatus.Running))
                {
                    this.cancelToken.Cancel();
                    if (wallclock.Elapsed > new TimeSpan(0, 0, 30))
                    {
                        break;
                    }
                }
                intersectionTask.Dispose();
                intersectionTask = null;
            }

            this.todoRayBuffers = new InputRayBufferCollection();
            this.doneRayBuffers = new OutputRayBufferCollection();
        }
        public bool lowLatency = false;
        public override void Trace(RayBuffer rayBuffer)
        {
            //if (rayBuffer.RaysInitialized == 0)
            //{
            //    Debugger.Break();
            //}
            var raysArg = rayBuffer.RaysInfo;
            var raysHandle = GCHandle.Alloc(raysArg, GCHandleType.Pinned);

            Rays.SetData(raysArg);
            uint size = lowLatency
                            ? (uint) (Math.Sqrt(RayBuffer.RayBufferSize / 8))
                            : (uint) (Math.Sqrt(RayBuffer.RayBufferSize));
            Context.Launch(0, size, size);
             RayHits.GetData(ref rayBuffer.rayHits);
            rayBuffer.traced = true;
            raysHandle.Free();
        }

        public override void Start()
        {
            if (started)
                return;
            started = true;
            if (wallclock == null)
            {
                wallclock = Stopwatch.StartNew();
            }
            else
            {
                wallclock.Start();
            }
            statsDeviceIdleTime = 0.0;
            statsDeviceTotalTime = 0.0;
            if (cancelToken == null)
            {
                cancelToken = new CancellationTokenSource();
                this.intersectionTask = Task.Factory.StartNew(() => RayIntersection(cancelToken.Token, this),
                    cancelToken.Token, TaskCreationOptions.LongRunning, TaskScheduler.Current);
            }
        }

        protected static void RayIntersection(CancellationToken ct, OptixIntersectionDevice device)
        {
#if VERBOSE
            try
            {
                
#endif
            Tracer.TraceLine("Intersection device {0}", device.GetType().Name);
            while (true)
            {
                if (ct.IsCancellationRequested || device.cancelToken.IsCancellationRequested)
                {
                    Tracer.TraceLine("Stopping intersection device {0}", device.GetType().Name);
                    break;
                }
                var t1 = 0;// device.wallclock.ElapsedTicks;
                Tuple<int, RayBuffer> rayBuffer = null;

#if OLD 
                if (!device.todoRayBuffers.TryDequeue(out rayBuffer))
                    continue;

#else
                rayBuffer = device.todoRayBuffers.Get();
                /*
                if (device.todoRayBuffers.IsEmpty)
                    continue;
                if (!device.todoRayBuffers.TryDequeue(out rayBuffer))
                {
                    continue;
                }
                 * */
#endif

                var t2 = DateTime.UtcNow.Ticks;
                device.Trace(rayBuffer.Item2);
                //Tracer.TraceLine("Result returns {0} items ", buffer.rayHits.Length);
                //            Tracer.TraceLine("Intersections {0} items ", buffer.rayHits.Count(p => p.index != uint.MaxValue));
                var t3 = device.wallclock.ElapsedTicks;

                device.statsDeviceIdleTime += t2 - t1;
                device.statsDeviceTotalTime += t3 - t1;
                device.statsTotalRayCount += rayBuffer.Item2.RaysInfo.Length;
                device.doneRayBuffers.Put(rayBuffer.Item1, rayBuffer.Item2);
                //device.doneRayBuffers[rayBuffer.Item1].Enqueue(rayBuffer.Item2);
            }

#if VERBOSE
            }
            catch (Exception ex)
            {
                Tracer.TraceLine("Intersection device exception");
                //Tracer.TraceLine("Error triangle {0}", currentTriangleIndex);
                Tracer.TraceLine(ex.Message);
                Tracer.TraceLine(ex.StackTrace);

            }
#endif
        }

        public override void Dispose()
        {
            if (Context != null)
            {
                Rays.Dispose();
                RayHits.Dispose();
                GeoGroup.Dispose();
                Context.Dispose();
            }
        }
    }
}
